## 使用deepseek学习milvus中的内容并能够进行问答功能

# 显示所有可用集合
# 选择要查询的文档集合
# 使用DeepSeek模型进行问答
# 显示回答和相关来源

from config import envConfig
from pymilvus import connections, utility
from langchain_ollama import OllamaEmbeddings
from langchain_community.vectorstores import Milvus
from langchain_community.llms import Ollama
from langchain.chains import RetrievalQA
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('qa_system')

EMBEDDING_MODEL = "deepseek-r1:1.5b"
LLM_MODEL = "deepseek-r1:1.5b"


def connect_milvus():
    """连接Milvus数据库"""
    try:
        connections.connect(
            alias="default",
            host=envConfig.MILVUS_HOST,
            port=envConfig.MILVUS_PORT,
            user=envConfig.MILVUS_USER,
            password=envConfig.MILVUS_PASSWORD,
            secure=False
        )
        return True
    except Exception as e:
        logger.error(f"Milvus连接失败: {str(e)}")
        return False


def list_collections():
    """列出所有集合"""
    if not connect_milvus():
        return []

    try:
        return utility.list_collections()
    except Exception as e:
        logger.error(f"获取集合列表失败: {str(e)}")
        return []


def initialize_qa_system(collection_name):
    """初始化问答系统"""
    # 1. 初始化嵌入模型
    embeddings = OllamaEmbeddings(
        model=EMBEDDING_MODEL,
        base_url=envConfig.OLLAMA_URL
    )

    # 2. 创建向量库连接
    vector_db = Milvus(
        embedding_function=embeddings,
        collection_name=collection_name,
        connection_args={
            "host": envConfig.MILVUS_HOST,
            "port": envConfig.MILVUS_PORT,
            "user": envConfig.MILVUS_USER,
            "password": envConfig.MILVUS_PASSWORD
        }
    )

    # 3. 初始化语言模型
    llm = Ollama(
        model=LLM_MODEL,
        base_url=envConfig.OLLAMA_URL,
        temperature=0.2,
        num_ctx=4096,
        system="你是一个专业的文档助手，根据提供的文档内容回答问题。"
    )

    # 4. 创建检索式QA链
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vector_db.as_retriever(search_kwargs={"k": 3}),
        return_source_documents=True,
        verbose=False
    )

    logger.info(f"问答系统初始化完成 (集合: {collection_name})")
    return qa_chain


def interactive_qa(qa_chain, collection_name):
    """交互式问答循环"""
    print(f"\n===== 文档问答系统 =====")
    print(f"当前文档集合: {collection_name}")
    print("输入问题开始查询，输入 'exit' 退出")

    while True:
        question = input("\n问题: ").strip()
        if not question:
            continue
        if question.lower() in ["exit", "quit"]:
            break

        try:
            # 执行查询
            result = qa_chain.invoke({"query": question})

            # 显示回答
            print("\n回答:", result["result"])

            # 显示来源
            if result["source_documents"]:
                print("\n来源:")
                for i, doc in enumerate(result["source_documents"], 1):
                    page_num = doc.metadata.get('page', 0) + 1
                    source = doc.metadata.get('source', '未知来源')
                    print(f"{i}. 文档: {source}, 页码: {page_num}")
                    print(f"   内容: {doc.page_content[:120]}...")
            else:
                print("\n未找到相关来源")

        except Exception as e:
            logger.error(f"问答过程中出错: {str(e)}")
            print("抱歉，处理问题时出错，请重试")


def main():
    """主函数：问答系统界面"""
    if not connect_milvus():
        return

    # 获取集合列表
    collections = list_collections()
    if not collections:
        print("Milvus中没有可用的集合")
        return

    # 选择集合
    print("\n可用的文档集合:")
    for i, col in enumerate(collections, 1):
        print(f"{i}. {col}")

    try:
        choice = int(input("请选择要查询的集合编号: ").strip())
        if 1 <= choice <= len(collections):
            collection_name = collections[choice - 1]
        else:
            print("无效的选择")
            return
    except ValueError:
        print("请输入有效的数字")
        return

    # 初始化问答系统
    qa_chain = initialize_qa_system(collection_name)

    # 进入问答循环
    interactive_qa(qa_chain, collection_name)


if __name__ == "__main__":
    main()